﻿using System;
using System.Collections.Generic;
using System.Text;

namespace Hu01
{
	public class HuffmanTree
	{
		short[] lookup;
		short[] left;
		short[] right;

		public HuffmanTree(byte[] codeLengths)
		{
			// Calculate count for each bit length
			uint[] bitLengthCount = new uint[16];
			for (int i = 0; i < 256; i++)
			{
				bitLengthCount[codeLengths[i] & 0xf]++;
				bitLengthCount[codeLengths[i] >> 4]++;
			}
			bitLengthCount[0] = 0;

			// Generate initial codes for each bit length (RFC 1951, section 3.2.2.)
			uint[] nextCode = new uint[16];
			uint tempCode = 0;
			for (int bits = 1; bits < 16; bits++)
			{
				tempCode = (tempCode + bitLengthCount[bits - 1]) << 1;
				nextCode[bits] = tempCode;
			}

			// Generate lookup table
			short avail = 0;
			left = new short[512];
			right = new short[512];
			lookup = new short[1024];
			for (short literal = 0; literal < 512; literal++)
			{
				// Each nibble in the codeLengths contains code length for one character
				byte length = (byte)((codeLengths[literal >> 1] >> ((literal & 1) << 2)) & 0xf);
				if (length != 0)
				{
					// Generate next code for particular code length
					short code = (short)nextCode[length]++;
					// Value for lookup table entry (decoded literal and code length)
					short tableValue = (short)((literal << 4) | length);

					if (length <= 10)
					{
						// Short code (<= 10 bits)
						for (int j = (1 << (10 - length)) - 1; j >= 0; j--)
							lookup[(code << (10 - length)) + j] = tableValue;
					}
					else
					{
						// Long code (> 10 bits), use tree for overflow bits
						int overflowBits = length - 10;
						int codeBitMask = 1 << (overflowBits - 1);
						int index = (code >> overflowBits) & 0x3ff;
						short[] array = lookup;
						do
						{
							short value = array[index];
							if (value == 0)
								array[index] = value = (short)-(avail++ + 1);
							array = ((code & codeBitMask) == 0) ? left : right;
							index = -value - 1;
							codeBitMask >>= 1;
							overflowBits--;
						} while (overflowBits != 0);
						array[index] = tableValue;
					}
				}
			}
		}

		public short Lookup(InputBuffer inputBuffer)
		{
			short tableValue = lookup[inputBuffer.PeekBits(10)];
			if (tableValue < 0)
			{
				// Long code (> 10 bits)
				inputBuffer.SkipBits(10);
				do
				{
					if (inputBuffer.PeekBits(1) > 0)
						tableValue = right[-tableValue - 1];
					else
						tableValue = left[-tableValue - 1];
					inputBuffer.SkipBits(1);
				} while (tableValue < 0);
			}
			else
			{
				inputBuffer.SkipBits((byte)(tableValue & 0xf));
			}
			return (short)(tableValue >> 4);
		}
	}
}
